from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import train_test_split

# 加载SMS垃圾短信数据集
with open('SMSSpamCollection', 'r', encoding='utf8') as f:
    sms = [line.split('\t') for line in f]
y, x = zip(*sms)
# 为了测试可以先取少量样本
# n = 100
# y, x = y[:n], x[:n]
print(len(x))
y = [1 if label == 'spam' else 0 for label in y]  # 标签为spam的样本的整数标签为1，表示是垃圾短信，反之0则不是垃圾短信
X_train, X_test, y_train, y_test = train_test_split(x, y)


# SMS垃圾短信数据集特征提取
counter = CountVectorizer(token_pattern='[a-zA-Z]{2,}', max_features=2000)  # token_pattern='[a-zA-Z]{2,}'指定了一个正则表达式，用于匹配由两个或两个以上连续英文字母（不区分大小写）组成的字符串序列。
X_train = counter.fit_transform(X_train)  # 得到的结果是一个sklearn的稀疏矩阵，不能直接用len()方法直接获取其长度（样本数量）
X_test = counter.transform(X_test)   # 稀疏矩阵的元素
X_train, X_test = X_train.toarray(), X_test.toarray()  # 转为numpy数组形式（主要这会使得转换后的变量占据更多空间）
print(X_train.shape[0],X_test.shape[0])


